import numpy as np
from plyfile import PlyData
import shutil
import os

def standardize_bbox(pcl, points_per_object):
    pt_indices = np.random.choice(
        pcl.shape[0], pcl.shape[0], replace=False)
    np.random.shuffle(pt_indices)
    pcl = pcl[pt_indices]  # n by 3
    mins = np.amin(pcl, axis=0)
    maxs = np.amax(pcl, axis=0)
    center = (mins + maxs) / 2.
    scale = np.amax(maxs - mins)
    #print("Center: {}, Scale: {}".format(center, scale))
    result = ((pcl - center) / scale).astype(np.float32)  # [-0.5, 0.5]
    return result


xml_head = \
    """
<scene version="0.5.0">
    <integrator type="path">
        <integer name="maxDepth" value="-1"/>
    </integrator>
    <sensor type="perspective">
        <float name="farClip" value="100"/>
        <float name="nearClip" value="0.1"/>
        <transform name="toWorld">
            <lookat origin="4,0,0" target="0,0,0" up="0,0,1"/>
        </transform>
        <float name="fov" value="25"/>

        <sampler type="ldsampler">
            <integer name="sampleCount" value="256"/>
        </sampler>
        <film type="ldrfilm">
            <integer name="width" value="1600"/>
            <integer name="height" value="1200"/>
            <rfilter type="gaussian"/>
            <boolean name="banner" value="false"/>
        </film>
    </sensor>

    <bsdf type="roughplastic" id="surfaceMaterial">
        <string name="distribution" value="ggx"/>
        <float name="alpha" value="0.05"/>
        <float name="intIOR" value="1.46"/>
        <rgb name="diffuseReflectance" value="1,1,1"/> <!-- default 0.5 -->
    </bsdf>

"""

xml_ball_segment = \
    """
    <shape type="sphere">
        <float name="radius" value="0.02"/>
        <transform name="toWorld">
            <translate x="{}" y="{}" z="{}"/>
        </transform>
        <bsdf type="diffuse">
            <rgb name="reflectance" value="{},{},{}"/>
        </bsdf>
    </shape>
"""

xml_tail = \
    """
    <shape type="rectangle">
        <ref name="bsdf" id="surfaceMaterial"/>
        <transform name="toWorld">
            <scale x="10" y="10" z="1"/>
            <translate x="0" y="0" z="-0.5"/>
        </transform>
    </shape>

    <shape type="rectangle">
        <transform name="toWorld">
            <scale x="10" y="10" z="1"/>
            <rotate y="1" angle="90"/>
            <translate x="-5" y="0" z="0"/>
        </transform>
        <emitter type="area">
            <rgb name="radiance" value="6,6,6"/>
        </emitter>
    </shape>
</scene>
"""


def colormap(x, y, z):
    vec = np.array([x, y, z])
    vec = np.clip(vec, 0.001, 1.0)
    norm = np.sqrt(np.sum(vec ** 2))
    vec /= norm
    return [0.3, 0.3, 0.3]
    # return [vec[0], vec[1], vec[2]]


xml_segments = [xml_head]

idxxx = 0

file_path = '/home/ljl/hdd/lidarcap_result_gallery_video/'
xml_file_path = '/home/ljl/hdd/lidarcap_video/lidarcap_result_gallery_video_cloud_xml/'
dirs = [filename for filename in os.listdir(file_path)]
for dir in dirs:
    file_dirs = [filename for filename in os.listdir(file_path + dir + '/')]
    for file_dir in file_dirs:

        filenames = [filename for filename in os.listdir(file_path + dir + '/' + file_dir)]

        if os.path.exists(xml_file_path + dir + '/' + file_dir + '/'):
            shutil.rmtree(xml_file_path + dir + '/' + file_dir + '/')
            os.makedirs(xml_file_path + dir + '/' + file_dir + '/' )
        else:
            os.makedirs(xml_file_path + dir + '/' + file_dir + '/')

        for filename in filenames:
            if filename.endswith('cloud.ply'):
                idxxx = idxxx + 1
                print(idxxx)
                file_index = filename[0:-4]
                in_filename = file_path + dir + '/' + file_dir + '/' + filename

                xmlname = xml_file_path + dir + '/' + file_dir + '/' + file_index + '.xml'

                pcl1 = PlyData.read(in_filename)['vertex'].data

                theta_x = 0
                theta_y = 0
                theta_z = 180 / 180 * np.pi

                Rot_x = np.identity(3)
                Rot_y = np.identity(3)
                Rot_z = np.identity(3)

                Rot_x[1][1] = np.cos(theta_x)
                Rot_x[1][2] = -np.sin(theta_x)
                Rot_x[2][1] = np.sin(theta_x)
                Rot_x[2][2] = np.cos(theta_x)

                Rot_y[0][0] = np.cos(theta_y)
                Rot_y[0][2] = np.sin(theta_y)
                Rot_y[2][0] = -np.sin(theta_y)
                Rot_y[2][2] = np.cos(theta_y)

                Rot_z[0][0] = np.cos(theta_z)
                Rot_z[0][1] = -np.sin(theta_z)
                Rot_z[1][0] = np.sin(theta_z)
                Rot_z[1][1] = np.cos(theta_z)

                pcl = np.array([[x, y, z] for x, y, z in pcl1])
                #print(pcl.shape)
                pcl = np.matmul(np.matmul(np.matmul(pcl, Rot_x.T), Rot_y.T), Rot_z.T)
                pcl = standardize_bbox(pcl, 2048)
                # pcl = pcl[:, [2, 0, 1]]
                # pcl[:, 0] *= -1
                pcl[:, 2] += 0.1

                for i in range(pcl.shape[0]):
                    color = colormap(pcl[i, 0] + 0.5, pcl[i, 1] +
                                     0.5, pcl[i, 2] + 0.5 - 0.0125)
                    xml_segments.append(xml_ball_segment.format(
                        pcl[i, 0], pcl[i, 1], pcl[i, 2], *color))
                xml_segments.append(xml_tail)

                xml_content = str.join('', xml_segments)

                with open(xmlname, 'w') as f:
                    f.write(xml_content)
